import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models.resnet import ResNet, Bottleneck

class UNetResNet50(nn.Module):
    def __init__(self, n_classes):
        super(UNetResNet50, self).__init__()

        # Load pretrained ResNet50 model
        resnet = models.resnet50(pretrained=True)

        # Extract layers from ResNet-50
        self.encoder0 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu)
        #定义编码器的第一部分，包括卷积、批量归一化和激活函数。
        self.encoder1 = nn.Sequential(resnet.maxpool, resnet.layer1)
        #定义编码器的第二部分，包括最大池化层和ResNet-50的第一个残差块。
        self.encoder2 = resnet.layer2
        #定义编码器的第三部分，包括ResNet-50的第二个残差块。
        self.encoder3 = resnet.layer3
        self.encoder4 = resnet.layer4

        # Decoder part (upsampling)
        self.upconv4 = self.up_conv(2048, 1024)
        #定义解码器的第一部分，包括一个上采样层，输入通道数为 2048，输出通道数为 1024。
        self.decoder4 = self.conv_block(2048, 1024)
#定义卷积块，输入通道数为 2048，输出通道数为 1024。
        self.upconv3 = self.up_conv(1024, 512)
        self.decoder3 = self.conv_block(1024, 512)

        self.upconv2 = self.up_conv(512, 256)
        self.decoder2 = self.conv_block(512, 256)

        self.upconv1 = self.up_conv(256, 64)
        self.decoder1 = self.conv_block(128, 64)

        self.conv_last = nn.Conv2d(64, n_classes, kernel_size=1)

        # Add upsampling layer to match target size
        self.upsample = nn.Upsample(size=(256, 256), mode='bilinear', align_corners=True)

    def conv_block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def up_conv(self, in_channels, out_channels):
        return nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)

    def forward(self, x):
        # Encoder
        x0 = self.encoder0(x)# 编码器部分，依次提取特征
        x1 = self.encoder1(x0)
        x2 = self.encoder2(x1)
        x3 = self.encoder3(x2)
        x4 = self.encoder4(x3)

        # Decoder with skip connections
        d4 = self.upconv4(x4)
        d4 = torch.cat([d4, x3], dim=1)
        d4 = self.decoder4(d4)

        d3 = self.upconv3(d4)
        d3 = torch.cat([d3, x2], dim=1)
        d3 = self.decoder3(d3)

        d2 = self.upconv2(d3)
        d2 = torch.cat([d2, x1], dim=1)
        d2 = self.decoder2(d2)

        d1 = self.upconv1(d2)
        d1 = torch.cat([d1, x0], dim=1)
        d1 = self.decoder1(d1)

        out = self.conv_last(d1)

        # Upsample output to match target size
        out = self.upsample(out)

        return out